import pandas as pd
import json
import uuid
import os

# 指定文件路径
nodes_file = "../Ti_data/nodes_info_5.csv"
edges_file = "../Ti_data/edges_info_5.csv"

# 解析文件名，提取后缀
nodes_suffix = os.path.splitext(os.path.basename(nodes_file))[0].replace("nodes_", "")
edges_suffix = os.path.splitext(os.path.basename(edges_file))[0].replace("edges_", "")

# 确保两者的后缀一致
output_suffix = nodes_suffix if nodes_suffix == edges_suffix else "output"

# 读取 CSV 文件
nodes_df = pd.read_csv(nodes_file)
edges_df = pd.read_csv(edges_file)

# 生成 node_id 到 UUID 的映射
node_id_map = {row["node_id"]: str(uuid.uuid4()) for _, row in nodes_df.iterrows()}

# 处理节点数据
nodes = []
for _, row in nodes_df.iterrows():
    node = {
        "id": node_id_map[row["node_id"]],  # 使用 UUID 作为 id
        "data": {key: row[key] for key in nodes_df.columns if key != "node_id" and pd.notna(row[key])}
    }
    nodes.append(node)

# 处理边数据，去重
edges_set = set()
edges = []
for _, row in edges_df.iterrows():
    source_uuid = node_id_map.get(row["source"], str(uuid.uuid4()))  # 可能存在无匹配的情况
    target_uuid = node_id_map.get(row["target"], str(uuid.uuid4()))

    data_dict = {key: row[key] for key in edges_df.columns if key not in ["source", "target"] and pd.notna(row[key])}
    edge_tuple = (source_uuid, target_uuid, json.dumps(data_dict, sort_keys=True))

    if edge_tuple not in edges_set:
        edges_set.add(edge_tuple)
        edge = {
            "id": str(uuid.uuid4()),  # 生成唯一的边 ID
            "source": source_uuid,
            "target": target_uuid,
            "data": data_dict
        }
        edges.append(edge)

# 组合成最终的 JSON 结构
output = {
    "nodes": nodes,
    "edges": edges
}

# 自动生成 JSON 输出文件名
output_file = f"output_{output_suffix}.json"
with open("../Ti_data/" + output_file, "w", encoding="utf-8") as f:
    json.dump(output, f, ensure_ascii=False, indent=2)

print(f"JSON 数据已生成：{output_file}")
